from classifier.classifier import Classifier
from sklearn import tree
import pydotplus
import os
import numpy as np
from BigDataWeb import config


# 决策树
class DecisionTree(Classifier):
    # 特征选择标准(gini,entropy)
    criterion = "gini"
    # 最大深度
    max_depth = None
    # 决策树图片路径
    graph_file_path = ""
    
    def __init__(self):
        Classifier.__init__(self)
        self.algorithm_name = "决策树"
        self.ipynb_template_name = "decision_tree-template.ipynb"
        self.graph_file_path = self.work_folder_path + "/graph.png"
        
    def setOutPutFieldName(self, output_field_name):
        # 分类字段必须是字符串
        self.data_source[output_field_name] = self.data_source[output_field_name].astype(np.str)
        Classifier.setOutPutFieldName(self, output_field_name)
        
    def implent(self): 
        Classifier.implent(self)
        # 构造模型
        self.algorithm = tree.DecisionTreeClassifier(criterion=self.criterion, max_depth=self.max_depth)
        # 训练模型
        self.algorithm.fit(self.train_inputs, self.train_outputs)
        # 评估模型
        self.score = self.algorithm.score(self.test_inputs, self.test_outputs)
        # 预测
        self.predict_output_values = self.algorithm.predict(self.predict_input_values)
        self.predict_output_probas = self.algorithm.predict_proba(self.predict_input_values)
        # 生成png图像
        # 添加graphviz路径至系统环境变量 
        os.environ["PATH"] += os.pathsep + config.graphviz_path
        # 将决策树导出为dot_data格式数据
        dot_data = tree.export_graphviz(self.algorithm, out_file=None,
                                feature_names=self.input_field_names,
                                class_names=self.algorithm.classes_,
                                filled=True) 
        # 规避中文出现乱码
        dot_data = dot_data.replace("node [shape=box, style=\"filled\", color=\"black\"] ;", "node [shape=box, style=\"filled\", color=\"black\",fontname=\"SimSun\"] ;")
        graph = pydotplus.graph_from_dot_data(dot_data)
        # 导出图片
        graph.write_png(self.graph_file_path)
    
    def prepareIpynbItems(self):
        Classifier.prepareIpynbItems(self)
        self.ipynb_items["#criterion#"] = self.criterion
        self.ipynb_items["#max_depth#"] = self.max_depth
